import lxml.etree as ET
from typing import List, Tuple
from log import logger

# logger = logging.getLogger(__name__)


def load_voc_annotation(xml_path: str) -> Tuple[List[List[int]], List[str]]:
    try:
        tree = ET.parse(xml_path)
        root = tree.getroot()
        bboxes, labels = [], []
        objects = root.findall("object")
        if not objects:
            logger.warning(f"No objects found in {xml_path}")
        for obj in objects:
            bbox = obj.find("bndbox")
            name = obj.find("name")
            if bbox is None or name is None:
                logger.error(f"Missing bndbox or name in object in {xml_path}")
                continue
            try:
                xmin = int(bbox.find("xmin").text)
                ymin = int(bbox.find("ymin").text)
                xmax = int(bbox.find("xmax").text)
                ymax = int(bbox.find("ymax").text)
                label = name.text
                if any(v is None for v in [xmin, ymin, xmax, ymax, label]):
                    logger.error(f"Invalid bbox or label in {xml_path}")
                    continue
                bboxes.append([xmin, ymin, xmax, ymax])
                labels.append(label)
            except (ValueError, AttributeError) as e:
                logger.error(f"Error parsing object in {xml_path}: {str(e)}")
                continue
        logger.info(f"Loaded {len(bboxes)} objects from {xml_path}")
        return bboxes, labels
    except ET.ParseError as e:
        logger.error(f"Failed to parse XML file {xml_path}: {str(e)}")
        raise


def save_voc_annotation(
    path: str, filename: str, image_shape: Tuple[int, int, int], bboxes: List[List[int]], labels: List[str]
) -> None:
    root = ET.Element("annotation")
    ET.SubElement(root, "filename").text = filename
    size = ET.SubElement(root, "size")
    ET.SubElement(size, "width").text = str(image_shape[1])
    ET.SubElement(size, "height").text = str(image_shape[0])
    depth = image_shape[2] if len(image_shape) > 2 else 3
    ET.SubElement(size, "depth").text = str(depth)

    for label, box in zip(labels, bboxes):
        obj = ET.SubElement(root, "object")
        ET.SubElement(obj, "name").text = label
        bndbox = ET.SubElement(obj, "bndbox")
        x1, y1, x2, y2 = map(int, box)
        ET.SubElement(bndbox, "xmin").text = str(x1)
        ET.SubElement(bndbox, "ymin").text = str(y1)
        ET.SubElement(bndbox, "xmax").text = str(x2)
        ET.SubElement(bndbox, "ymax").text = str(y2)

    ET.ElementTree(root).write(path, pretty_print=True)
